/*
 * hmm_test.cc
 *
 *  Created on: Sep 8, 2011
 *      Author: tqlong
 */

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <stack>

#include "../dhmm.h"
#include "../dur_hmm.h"
#include "../dur_hmm_markov.h"
#include "../position_model.h"
#include "../anneal.h"
#include "stl-helper.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void classic_hmm_test()
{
  std::cout << " ================ CLASSIC HMM ================ " << std::endl;
  typedef DiscreteHMM hmm_t;
  typedef DiscreteHMM::joint_seq_t seq_t;
  typedef DiscreteHMM::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();

  seq_t seq;
  hmm.generate(50, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 100;
  int T = 500;
  vseq_t data;
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    hmm.generate(T, seq);
    data.push_back(seq.oseq);
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.print_params();

  trained_hmm.train(data);
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();
}

void duration_hmm_test()
{
  std::cout << " ================ Duration HMM ================ " << std::endl;
  typedef DurationHMM hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();
  hmm.save(vm["output"].as<std::string>().c_str());

  seq_t seq;
  hmm.generate(10, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 100;
  int T = 200;
  vseq_t data;
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    hmm.generate(T, seq);
    data.push_back(seq.oseq);
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data);
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();
}

void duration_hmm_compare_test()
{
  std::cout << " ================ Duration HMM convergence ================ " << std::endl;
  typedef DurationHMM hmm_t;

  hmm_t hmm;
  hmm.load(vm["output"].as<std::string>().c_str());
  for (int iter = 0; iter < 40; iter++)
  {
    std::stringstream s;
    s << vm["prefix"].as<std::string>().c_str() << iter << ".txt";

    hmm_t iter_hmm;
    iter_hmm.load(s.str().c_str());
    double diff = hmm.diff(iter_hmm);
    std::cout << "iter=" << iter << " diff=" << diff << std::endl;

    std::stringstream sf;
    sf << vm["prefix"].as<std::string>().c_str() << "diff.txt";
    std::ofstream llfile(sf.str().c_str(), std::ios::app);
    llfile << iter << " " << diff << std::endl;
    llfile.close();
  }
}

void duration_hmm_markov_test()
{
  std::cout << " ================ Duration HMM ================ " << std::endl;
  typedef DurationHMMMarkov hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();
  hmm.save(vm["output"].as<std::string>().c_str());

  seq_t seq;
  hmm.generate(10, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 100;
  int T = 200;
  vseq_t data;
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    hmm.generate(T, seq);
    data.push_back(seq.oseq);
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data);
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();
}

void duration_hmm_markov_compare_test()
{
  std::cout << " ================ Duration HMM convergence ================ " << std::endl;
  typedef DurationHMMMarkov hmm_t;

  hmm_t hmm;
  hmm.load(vm["output"].as<std::string>().c_str());
  for (int iter = 0; iter < 50; iter++)
  {
    std::stringstream s;
    s << vm["prefix"].as<std::string>().c_str() << iter << ".txt";

    hmm_t iter_hmm;
    iter_hmm.load(s.str().c_str());
    double diff = hmm.diff(iter_hmm);
    std::cout << "iter=" << iter << " diff=" << diff << std::endl;

    std::stringstream sf;
    sf << vm["prefix"].as<std::string>().c_str() << "diff.txt";
    std::ofstream llfile(sf.str().c_str(), std::ios::app);
    llfile << iter << " " << diff << std::endl;
    llfile.close();
  }
}

void motif_test()
{
  std::cout << " ================ Motif with Duration HMM ================ " << std::endl;
  typedef DurationHMM hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();

  seq_t seq;
  hmm.generate(8, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 1000;
  int T = 8;
  vseq_t data;
  std::ofstream output(vm["seqfile"].as<std::string>().c_str());
  const char *atcg = "ATCG";
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    while (seq.oseq.size() < 11) 
      hmm.generate(T, seq);
    data.push_back(seq.oseq);
    output << ">" << n+1  << std::endl;
    for (unsigned int i = 0; i < seq.oseq.size()-1; i++) {
      output << atcg[seq.oseq[i]];
    }
    output << std::endl;
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.trans_ = hmm.trans_; // fix the transmission probability table
  trained_hmm.start_ = hmm.start_; // and the starting probabilities
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data,  vm["iter"].as<int>());
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();

  {
    std::vector<int> pos(data.size());

    for (unsigned int i = 0; i < data.size(); i++) {
    
      seq_t map_seq(data[i]);
      trained_hmm.map(map_seq);
      //std::cout << "MAP sequence: " << std::endl;
      //map_seq.print();
  
      pos[i] = -1;
      for (unsigned int t = 0; t < map_seq.qseq.size(); t++) {
	if (map_seq.qseq[t] == 1 && map_seq.qseq[t+5] == 6) pos[i] = t;
      }
    }

    std::ofstream output(vm["output"].as<std::string>().c_str());
    for (unsigned int i = 0; i < data.size(); i++) {
      output << ">" << i+1 << " pos = " << pos[i] << std::endl;
      if (pos[i] != -1) {
	for (int k = 0; k < 20-pos[i]; k++) output << " ";
	for (int k = 0; k < pos[i]; k++) output << atcg[data[i][k]];
	output << " ";
	for (int k = pos[i]; k < pos[i]+5; k++) output << atcg[data[i][k]];
	output << " ";
	for (unsigned int k = pos[i]+5; k < data[i].size()-1; k++) output << atcg[data[i][k]];
	output << std::endl;
      }
      else {
	for (unsigned int k = 0; k < data[i].size(); k++) output << atcg[data[i][k]];
	output << std::endl;
      }
    }
    output.close();
  }
}

template<typename vseq_t>
void load_sequence(const char* fn, vseq_t& data, const char* atcg = "ATCG")
{
  std::ifstream input(fn);
  //int i = 1;
  while (!input.eof()) {
    std::string s;
    getline(input, s);
    //const char *atcg = "ATCG";
    if (s.length() > 0 && strchr(atcg, s[0])) {
      //std::cout << i++ << " " << s  << std::endl;
      typename vseq_t::value_type seq(s.length()+1);
      for (unsigned int i = 0; i < s.length(); i++) {
	const char* pos = strchr(atcg, s[i]);
	if (!pos) { std::cerr << "error unknown char " << s[i] << std::endl; exit(1); }
	seq[i] = pos - atcg;
      }
      seq[s.length()] = 4;
      data.push_back(seq);
    }
  }
}

void real_motif_test()
{
  std::cout << " ================ Real Motif with Duration HMM ================ " << std::endl;
  typedef DurationHMM hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());

  // int N = 1000;
  // int T = 8;
  vseq_t data;
  
  load_sequence(vm["seqfile"].as<std::string>().c_str(), data);

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.trans_ = hmm.trans_; // fix the transmission probability table
  trained_hmm.start_ = hmm.start_; // and the starting probabilities
  trained_hmm.dprob_ = hmm.dprob_;
  for (int j = 0; j < trained_hmm.n_state(); j++) {
    trained_hmm.emis_(trained_hmm.n_observation()-1, j) = 0.0;
    if (j == trained_hmm.n_state()-1) {
      trained_hmm.emis_(trained_hmm.n_observation()-1, j) = 1.0;
      for (int k = 0; k < trained_hmm.n_observation()-1; k++)
	trained_hmm.emis_(k, j) = 0.0;
    }
    trained_hmm.emis_.col(j) /= arma::sum(trained_hmm.emis_.col(j));
  }
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data, vm["iter"].as<int>());
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();

  std::vector<int> pos(data.size());
  for (unsigned int i = 0; i < data.size(); i++) {
    
    seq_t map_seq(data[i]);
    trained_hmm.map(map_seq);
    //std::cout << "MAP sequence: " << std::endl;
    //map_seq.print();
  
    pos[i] = -1;
    for (unsigned int t = 0; t < map_seq.qseq.size(); t++) {
      if (map_seq.qseq[t] == 1 && map_seq.qseq[t+5] == 6) pos[i] = t;
    }
  }

  std::ofstream output(vm["output"].as<std::string>().c_str());
  const char *atcg = "ATCG";
  for (unsigned int i = 0; i < data.size(); i++) {
    output << ">" << i+1 << " pos = " << pos[i] << std::endl;
    if (pos[i] != -1) {
      for (int k = 0; k < 50-pos[i]; k++) output << " ";
      for (int k = 0; k < pos[i]; k++) output << atcg[data[i][k]];
      output << " ";
      for (int k = pos[i]; k < pos[i]+6; k++) output << atcg[data[i][k]];
      output << " ";
      for (unsigned int k = pos[i]+6; k < data[i].size()-1; k++) output << atcg[data[i][k]];
      output << std::endl;
    }
    else {
      for (unsigned int k = 0; k < data[i].size(); k++) output << atcg[data[i][k]];
      output << std::endl;
    }
  }
  output.close();
}

void position_test()
{
  std::cout << " ================ Real Motif with Duration HMM ================ " << std::endl;
  typedef PositionModel pm_t;
  typedef pm_t::vseq_t vseq_t;

  std::string outparam, outputfile, seqfile = vm["seqfile"].as<std::string>();
  if (vm.count("outparam"))
    outparam = vm["outparam"].as<std::string>();
  else
    outparam = seqfile + ".param";

  if (vm.count("output"))
    outputfile = vm["output"].as<std::string>();
  else
    outputfile = seqfile + ".output";

  vseq_t data;
  int motif_length = vm["G"].as<int>();
  int n_observation = 4;
  int max_seq_length = 30;
  load_sequence(seqfile.c_str(), data);
  for (unsigned int i = 0; i < data.size(); i++) {
    data[i].pop_back();
    if ((int) data[i].size() > max_seq_length)
      max_seq_length = data[i].size();
  }

  srand(time(NULL));
  pm_t pm(motif_length, n_observation, max_seq_length);
  pm.print_params();

  pm.train(data, vm["iter"].as<int>());
  std::cout << " -- Trained --" << std::endl;
  pm.print_params();

  std::vector<int> pos(data.size());
  for (unsigned int i = 0; i < data.size(); i++)
    pos[i] = pm.map_pos(data[i]);

  const char *atcg = "ATCG";
  std::ofstream output(outputfile.c_str());
  for (unsigned int i = 0; i < data.size(); i++) {
    output << ">" << i+1 << " pos = " << pos[i] << std::endl;
    if (pos[i] != -1) {
      for (int k = 0; k < 50-pos[i]; k++) output << " ";
      for (int k = 0; k < pos[i]; k++) output << atcg[data[i][k]];
      output << " ";
      for (int k = pos[i]; k < pos[i]+motif_length; k++) output << atcg[data[i][k]];
      output << " ";
      for (unsigned int k = pos[i]+motif_length; k < data[i].size(); k++) output << atcg[data[i][k]];
      output << std::endl;
    }
    else {
      for (unsigned int k = 0; k < data[i].size(); k++) output << atcg[data[i][k]];
      output << std::endl;
    }
  }
  output.close();
  
  pm.save_param(outparam);
}

void generate_motif() 
{
  int n_base = 4;
  int motif_length = vm["G"].as<int>();
  vvd_t freq;
  vvi_t data;
  vi_t motif_pos;
  int n_seq = vm["N"].as<int>();
  int seq_length = vm["L"].as<int>();
  vd_t pos_prob = arma::zeros<vd_t>(seq_length - motif_length);
  std::string seqFile = vm["seqfile"].as<std::string>();

  double m = vm["mean"].as<double>();
  double s = vm["std"].as<double>();
  double total = 0;
  for (int i = 0; i < sz(pos_prob); i++) {
    pos_prob[i] = exp(-0.5*sqr(i-m)/sqr(s));
    total += pos_prob[i];
  }

  Anneal a(n_base, motif_length);
  a.generate_freq(freq);
  a.generate_data(n_seq, seq_length, freq, pos_prob, data, motif_pos);
  a.save_data(seqFile, data, motif_pos);
  a.save_freq(seqFile+".truth", freq);

  std::string truthfile = seqFile + ".truth";
  std::ofstream truth(truthfile.c_str(), std::ios::app);
  for (int i = 0; i < sz(pos_prob); i++)
    truth << i << ",\t" << pos_prob[i]/total << ENDL;
  truth.close();
}

int extract_pos(const std::string& s) 
{
  std::stack<char> stack;
  for (int i = s.length()-1; i >= 0; i--) {
    if (s[i] >= '0' && s[i] <= '9') stack.push(s[i]);
    else break;
  }
  std::string ss;
  while (!stack.empty()) {
    ss.pb(stack.top());
    stack.pop();
  }
  std::stringstream sss(ss);
  int pos = -1;
  sss >> pos;
  //std::cout << "pos=" << pos << ENDL;
  return pos;
}

void parseFileToDoubles(const std::string& fileName, std::vector<double>& numbers)
{
  numbers.clear();
  std::ifstream file(fileName.c_str());
  std::string tmp;
  int count = 0;
  while (!file.eof()) {
    std::string line;
    std::getline(file, line);
    std::stringstream linestream(line);
    while (std::getline(linestream, tmp, ',')) {
      std::stringstream parser(tmp);
      double a;
      parser >> a;
      //std::cout << count++ << " " << a << ENDL;
      numbers.pb(a);
    }
  }
  file.close();
}

void compare()
{
  std::string outparam, outputfile, seqfile = vm["seqfile"].as<std::string>();
  if (vm.count("outparam"))
    outparam = vm["outparam"].as<std::string>();
  else
    outparam = seqfile + ".param";

  if (vm.count("output"))
    outputfile = vm["output"].as<std::string>();
  else
    outputfile = seqfile + ".output";

  std::string seqFile = seqfile, truthFile = seqfile + ".truth";
  std::string outFile = outputfile, outParam = outparam;
  
  std::ifstream seq(seqFile.c_str());
  std::ifstream out(outFile.c_str());

  double n_line = 0, acc = 0;
  while (!seq.eof()) {
    std::string seq_line, out_line;
    std::getline(seq, seq_line);    
    std::getline(out, out_line);

    //std::cout << "seq: " << seq_line << ENDL;
    //std::cout << "out: " << out_line << ENDL;
    if (seq_line.empty()) continue;

    int truth = extract_pos(seq_line);
    int predict = extract_pos(out_line);
    if (predict == truth) acc += 1;
    n_line += 1;

    std::getline(seq, seq_line);    
    std::getline(out, out_line);
  }
  std::vector<double> truthNumbers, paramNumbers;
  parseFileToDoubles(truthFile, truthNumbers);
  parseFileToDoubles(outParam, paramNumbers);
  
  // if (truthNumbers.size() != paramNumbers.size()) {
  //   FATAL_ERROR("number of doubles not equal");
  // }

  double l2 = 0;
  int N = truthNumbers.size();
  for (int i = 0; i < N && i < (int) paramNumbers.size(); i++) 
    l2 += sqr(truthNumbers[i] - paramNumbers[i]);

  std::cout << outFile << "\tacc = \t" << acc / n_line * 100 << "\t"
            << "l2 = \t" << sqrt(l2) << ENDL;
}

int main(int argc, char** argv)
{
  std::cerr << "hmm tests ..." << std::endl;
  srand(time(NULL));

  process_options(argc, argv);

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if (vm.count("help"))
    usage();
  else
  {
    if (vm.count("classic"))
      classic_hmm_test();
    else if (vm.count("duration"))
      duration_hmm_test();
    //    else if (vm.count("compare"))
    //  duration_hmm_compare_test();
    else if (vm.count("duration_markov"))
      duration_hmm_markov_test();
    else if (vm.count("compare_markov"))
      duration_hmm_markov_compare_test();
    else if (vm.count("motif"))
      motif_test();
    else if (vm.count("real"))
      real_motif_test();
    else if (vm.count("position"))
      position_test();
    else if (vm.count("anneal"))
      anneal_real_test();
    else if (vm.count("gen"))
      generate_motif();
    else if (vm.count("compare"))
      compare();
  }

  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  std::cerr << "Duration = " << duration << "\n";
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("hmmfile", po::value<std::string>()->default_value("hmm1.txt"), "input file (hmm)")
    ("seqfile", po::value<std::string>(), "sequence file")
    ("outparam", po::value<std::string>(), "output parameter file")
    ("output", po::value<std::string>(), "output file")
    ("prefix", po::value<std::string>()->default_value("output/"), "output prefix")
    ("iter", po::value<int>()->default_value(10), "number of iterations")    
//("fin", po::value<std::string>(&fin)->default_value(""), "input file (factor functions)")
    //("out", po::value<std::string>(&out)->default_value(""), "output file")
    ("classic", "test the classic discrete HMM")
    ("compare", "convergence rate")
    ("duration", "test the duration discrete HMM (independent)")
    ("compare_markov", "convergence rate (markov)")
    ("duration_markov", "test the duration discrete HMM (markov)")
    ("motif", "motif finding")
    ("real", "real motif finding")
    ("position", "position model for motif finding")
    ("anneal", "simulated annealing for motif finding")
    ("gen", "generate motif")
    ("compare", "compare")
    ("N", po::value<int>()->default_value(1000), "number of sequences")
    ("L", po::value<int>()->default_value(30), "sequence length")
    ("G", po::value<int>()->default_value(6), "motif length")
    ("mean", po::value<double>()->default_value(15), "mean of position")
    ("std", po::value<double>()->default_value(5), "standard deviation of position")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

